import cv2
import numpy as np


def extract_and_isolate_color(image_path, output_path, color):
    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        print("Error: Could not read the image.")
        return

    # 将图片转换为HSV颜色空间
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 定义颜色的HSV范围
    if color == "blue":
        lower_bound = np.array([100, 150, 0])
        upper_bound = np.array([140, 255, 255])
    elif color == "red":
        lower_bound1 = np.array([0, 120, 70])
        upper_bound1 = np.array([10, 255, 255])
        lower_bound2 = np.array([170, 120, 70])
        upper_bound2 = np.array([180, 255, 255])
        mask1 = cv2.inRange(hsv_image, lower_bound1, upper_bound1)
        mask2 = cv2.inRange(hsv_image, lower_bound2, upper_bound2)
        mask = cv2.bitwise_or(mask1, mask2)
    elif color == "yellow":
        lower_bound = np.array([20, 100, 100])
        upper_bound = np.array([30, 255, 255])
    elif color == "green":
        lower_bound = np.array([35, 100, 100])
        upper_bound = np.array([85, 255, 255])
    else:
        print("Error: Unsupported color.")
        return

    if color != "red":
        # 创建掩码
        mask = cv2.inRange(hsv_image, lower_bound, upper_bound)

    # 创建一个全白的图像
    white_background = np.full_like(image, 255)

    # 将指定颜色部分从原始图像中分离出来
    color_extracted = cv2.bitwise_and(image, image, mask=mask)

    # 反转掩码
    mask_inv = cv2.bitwise_not(mask)

    # 将非指定颜色部分变为白色
    non_color_part = cv2.bitwise_and(white_background, white_background, mask=mask_inv)

    # 合并指定颜色部分和白色背景
    result = cv2.add(color_extracted, non_color_part)

    # 保存结果图片
    cv2.imwrite(output_path, result)
    print(f"{color.capitalize()} color isolated and saved to {output_path}")
